package com.core.websocket;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import java.io.IOException;
import java.util.Hashtable;
import java.util.Map;

/**
 * Created with IntelliJ IDEA.
 * User: xf
 * Date: 2016/5/6
 * Time: 23:47
 * To change this template use Editor | File and Code Templates.
 */
public class MarcoHandler extends AbstractWebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(MarcoHandler.class);
    private static Map<String, WebSocketSession> sessionMap = new Hashtable<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        logger.debug("WebSocket Connection established!");
        addToSessionMap(session);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        logger.debug("Receiced Client Message: " + message.getPayload());
        broadcastAll(message.getPayload());
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        super.handleBinaryMessage(session, message);
    }

    @Override
    protected void handlePongMessage(WebSocketSession session, PongMessage message) throws Exception {
        super.handlePongMessage(session, message);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        removeFromSessionMap(session);
        logger.debug("Connection closed. Status: " + status);
    }

    private synchronized void addToSessionMap(WebSocketSession session) {
        if (!sessionMap.containsKey(session.getId())) {
            sessionMap.put(session.getId(), session);
        }
    }

    private void removeFromSessionMap(WebSocketSession session) {
        if (sessionMap.containsKey(session.getId())) {
            sessionMap.remove(session.getId());
        }
    }

    public void broadcastAll(String messaage) {
        for (Map.Entry<String, WebSocketSession> entry : sessionMap.entrySet()) {
            try {
                entry.getValue().sendMessage(new TextMessage(messaage));
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }
}
